
%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Policy function for assests and consumption
% given income shocks matrix and parameters of U function 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [AT,CT] = PoliFunctions(Ygrid,agrid3,T,gamma,beta,r,CAPH)

n = size(CAPH,1);  
J = length(agrid3);
ygrid =  Ygrid(:,:,T);

VT = zeros(length(ygrid),length(agrid3),T);
AT = zeros(length(ygrid),length(agrid3),T);
CT = zeros(length(ygrid),length(agrid3),T);

% last period:

% Imposed lowest asset accumulation for tomorrow, given that I'll die with prob 1
% AT are the optimal assets to pick by the end of T
% regardless of the assest I bring to T (columns)
% regardless of the income I end up in T (rows)

AT(:,:,T) = ones(n,J)*agrid3(1,1);

% consumption given assets I bring to T 
% and income in T
CT(:,:,T) = repmat(ygrid',1,J)+(1+r)*repmat(agrid3,n,1);
% value of consumption 
VT(:,:,T) = -1/gamma*exp(-gamma*CT(:,:,T));

% Solve Finite Problem

%C_W(t,y_t,a_t,a_t+1) = y_{T-1}+A_{T-1}*(1+r)-A_{T}



for k=1:(T-1)
    t = T-k;
    ygrid =  Ygrid(:,:,t);
    C_W = repmat(ygrid',1,J,J)+(1+r)*repmat(agrid3,n,1,J)-permute(repmat(agrid3,n,1,J),[1,3,2]); 
    CAPHt = CAPH(:,:,t);
    % given current y and FUTURE a
    E_V = CAPHt*VT(:,:,t+1);
    E_V = permute(repmat(E_V,1,1,J),[1,3,2]);
    W = -1/gamma*exp(-gamma*C_W)+beta*E_V;
    [V,I] = max(W,[],3);
    VT(:,:,t) = V;
    % assets I choose to bring to period t+1
    % given assets I have in period t and income in t
    % is the asset where the third dimension is I
    AT(:,:,t)=agrid3(I);
    CT(:,:,t) =repmat(ygrid',1,J)+(1+r)*repmat(agrid3,n,1)-AT(:,:,t);
end

end

